# coding: UTF-8
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn import metrics
from tensorboardX import SummaryWriter

from utils.utils import get_time_dif

# 权重初始化，默认xavier
def init_network(model, method='xavier', exclude='embedding', seed=1):
    for name, w in model.named_parameters():
        if exclude not in name:
            if 'weight' in name:
                if method == 'xavier':
                    nn.init.xavier_normal_(w)
                elif method == 'kaiming':
                    nn.init.kaiming_normal_(w)
                else:
                    nn.init.normal_(w)
            elif 'bias' in name:
                nn.init.constant_(w, 0)
            else:
                pass

def train(config, model, train_iter, val_iter, test_iter, loss_function=F.cross_entropy):
    start_time = time.time()
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)

    # 学习率指数衰减，每次epoch：学习率 = gamma * 学习率
    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)
    total_batch = 0  # 记录进行到多少batch
    val_best_loss = float('inf')
    writer = SummaryWriter(log_dir=config.log_path + time.strftime('%m-%d_%H.%M', time.localtime()))
    for epoch in range(config.num_epochs):
        print('Epoch [{}/{}]'.format(epoch + 1, config.num_epochs))
        train_loss =0 # 一个epoch的训练平均loss
        train_acc = 0 # 一个epoch的训练平均acc
        cnt =0 
        for i, (trains, labels, _) in enumerate(train_iter):
            outputs = model(trains)
            model.zero_grad()
            labels = labels.flatten()
            loss = loss_function(outputs, labels)
            train_loss += loss.item()
            loss.backward()
            optimizer.step()
                
            total_batch += 1

            true = labels.data.cpu()
            predic = torch.max(outputs.data, 1)[1].cpu()
            train_acc += metrics.accuracy_score(true, predic)

            cnt = i + 1

        # 每完成一个epoch输出在训练集和验证集上的效果
        train_loss /= cnt
        train_acc /= cnt
        val_acc, val_loss = evaluate(config, model, val_iter, loss_function)
        if val_loss < val_best_loss:
            val_best_loss = val_loss
            torch.save(model.state_dict(), config.save_path)
            improve = '*'
        else:
            improve = ''
        time_dif = get_time_dif(start_time)
        msg = 'Train Loss: {0:>5.2},  Train Acc: {1:>6.2%},  Val Loss: {2:>5.2},  Val Acc: {3:>6.2%},  Time: {4} {5}'
        print(msg.format(train_loss, train_acc, val_loss, val_acc, time_dif, improve))
        writer.add_scalar('loss/train', train_loss, epoch)
        writer.add_scalar('loss/val', val_loss, epoch)
        writer.add_scalar('acc/train', train_acc, epoch)
        writer.add_scalar('acc/val', val_acc, epoch)
        model.train()
        scheduler.step() # 学习率衰减
    writer.close()
    test(config, model, test_iter, loss_function)

def test(config, model, test_iter, loss_function=F.cross_entropy):
    if not(isinstance(model, list)):
        models = [model, ]
        configs = [config, ]
    else:
        models = model
        configs = config
    for i, model in enumerate(models):
        #model.load_state_dict(torch.load(configs[i].save_path))
        model.load_state_dict(torch.load(configs[i].save_path, map_location=config.device))
        model.eval()
    start_time = time.time()
    test_acc, test_loss, test_report, test_confusion = evaluate(configs, models, test_iter, loss_function, test=True)
    msg = 'Test Loss: {0:>5.2},  Test Acc: {1:>6.2%}'
    print(msg.format(test_loss, test_acc))
    print("Precision, Recall and F1-Score...")
    print(test_report)
    print("Confusion Matrix...")
    print(test_confusion)
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)

def evaluate(config, model, data_iter, loss_function, test=False):
    if not(isinstance(model, list)):
        models = [model, ]
        configs = [config, ]
    else:
        models = model
        configs = config
    for model in models:
        model.eval()
    loss_total = 0
    predict_all = np.array([], dtype=int)
    labels_all = np.array([], dtype=int)
    with torch.no_grad():
        for texts, labels, source in data_iter:
            outputs = 0
            for model in models:
                outputs += model(texts)
                labels = labels.flatten()
                loss = loss_function(outputs, labels)
                loss_total += loss
            labels = labels.data.cpu().numpy()
            predic = torch.max(outputs.data, 1)[1].cpu().numpy()   
            labels_all = np.append(labels_all, labels)
            '''
            for i, label in enumerate(labels):
                if label != predic[i]:
                    print('ground truth:', label)
                    print('predic:', predic[i])
                    print('source:', source[i])
                    print('- - -')
                    #input()
            '''
            predict_all = np.append(predict_all, predic)

    acc = metrics.accuracy_score(labels_all, predict_all)
    if test:
        report = metrics.classification_report(labels_all, predict_all, digits=4)
        confusion = metrics.confusion_matrix(labels_all, predict_all)
        return acc, loss_total / len(data_iter), report, confusion
    return acc, loss_total / len(data_iter)

if __name__ == "__main__":
    print('a')
    a = torch.FloatTensor([[-0.3867,  0.0470, -0.2407,  0.4931],
        [-0.3756,  0.5491, -0.0152,  0.5762],
        [-0.4836,  0.4270,  0.5146,  0.2515],
        [-0.3064,  0.2886, -0.1370,  0.4034],
        [-0.5047,  0.0930,  0.0552,  0.4332]])
    b = torch.LongTensor([0, 1, 3, 3, 0])
    '''
    F_loss = FocalLoss(4)
    loss = F_loss(a,b)
    '''